package com.adam.example.aggregation.bucket;

import co.elastic.clients.elasticsearch.ElasticsearchClient;
import co.elastic.clients.elasticsearch._types.aggregations.Aggregate;
import co.elastic.clients.elasticsearch._types.aggregations.Buckets;
import co.elastic.clients.elasticsearch._types.aggregations.FiltersAggregate;
import co.elastic.clients.elasticsearch._types.aggregations.FiltersBucket;
import co.elastic.clients.elasticsearch._types.query_dsl.Query;
import co.elastic.clients.elasticsearch.core.SearchResponse;
import com.adam.example.ElasticsearchClientHolder;

import java.io.IOException;
import java.util.HashMap;
import java.util.Map;

public class FiltersAggregationExample {

    public static void main(String[] args) throws IOException {
        ElasticsearchClient elasticsearchClient = ElasticsearchClientHolder.getINSTANCE().initializeAndGetClient();

        Map<String, Query> queryMap = new HashMap<>();
        queryMap.put("java", Query.of(queryBuilder -> queryBuilder
                .term(termQueryBuilder -> termQueryBuilder
                        .field("title").value("java")
                )
        ));
        queryMap.put("python", Query.of(queryBuilder -> queryBuilder
                .term(termQueryBuilder -> termQueryBuilder
                        .field("title").value("python")
                )
        ));
        SearchResponse<Void> searchResponse = elasticsearchClient.search(builder -> builder
                        .index("test_books")
                        .size(0)
                        .aggregations("agg", aggregationBuilder -> aggregationBuilder
                                .filters(filtersAggregationBuilder -> filtersAggregationBuilder
                                        .filters(queryBucketsBuilder -> queryBucketsBuilder
                                                .keyed(queryMap)
                                        )
                                )
                        )
                , Void.class);
        Aggregate aggregate = searchResponse.aggregations().get("agg");
        FiltersAggregate filtersAggregate = aggregate.filters();
        Buckets<FiltersBucket> buckets = filtersAggregate.buckets();
        Map<String, FiltersBucket> bucketMap = buckets.keyed();
        for(Map.Entry<String, FiltersBucket> entry: bucketMap.entrySet()) {
            String key = entry.getKey();
            FiltersBucket bucket = entry.getValue();
            System.out.println(key + ":" + bucket.docCount());
        }

        ElasticsearchClientHolder.getINSTANCE().close();
    }

}
